from __future__ import print_function
from burp import IBurpExtender, ITab, IHttpListener
from javax.swing import JSplitPane, JScrollPane, JTable, JPanel, JTextField, JButton, JLabel, Box
from javax.swing.table import DefaultTableModel, TableRowSorter
from java.awt import BorderLayout, FlowLayout, Font
from javax.swing.event import DocumentListener
from java.util import ArrayList
from java.net import URL
from java.io import File
from javax.swing import JFileChooser
from javax.swing import SwingUtilities
from java.awt.datatransfer import StringSelection
from java.awt import Toolkit
import codecs
import sys
from java.io import PrintWriter

def get_main_domain(hostname):
    parts = hostname.split('.')
    if len(parts) > 2:
        if parts[-2] in ('co', 'com', 'org', 'net', 'gov', 'edu', 'cn') and len(parts) > 2:
            return '.'.join(parts[-3:])
        else:
            return '.'.join(parts[-2:])
    return hostname

class BurpExtender(IBurpExtender, ITab, IHttpListener):

    def registerExtenderCallbacks(self, callbacks):
        self._callbacks = callbacks
        self._helpers = callbacks.getHelpers()
        
        sys.stdout = PrintWriter(callbacks.getStdout(), True)
        sys.stderr = PrintWriter(callbacks.getStderr(), True)
        
        self._callbacks.setExtensionName("Lazy Route")

        self.log = ArrayList()
        self.main_panel = JSplitPane(JSplitPane.VERTICAL_SPLIT)
        self.ui_font = Font("SansSerif", Font.PLAIN, 12)

        self.setup_ui()
        self._callbacks.addSuiteTab(self)
        self._callbacks.registerHttpListener(self)
        print("Lazy Route extension loaded successfully. All URLs will be processed.")
        return

    def getTabCaption(self):
        return "Lazy Route"

    def getUiComponent(self):
        return self.main_panel

    def setup_ui(self):
        top_panel = JPanel(FlowLayout(FlowLayout.LEFT))
        
        filter_label = JLabel("Filter (by Main/Sub Domain):")
        self.filter_field = JTextField(30)
        self.filter_field.getDocument().addDocumentListener(FilterListener(self))

        export_button = JButton("Export as CSV", actionPerformed=self.export_csv)
        copy_l1_button = JButton("Copy L1 Routes", actionPerformed=lambda e: self.copy_routes(3))
        copy_l2_button = JButton("Copy L2 Routes", actionPerformed=lambda e: self.copy_routes(4))
        clear_button = JButton(u"清除结果", actionPerformed=self.clear_results)

        for component in [filter_label, self.filter_field, export_button, copy_l1_button, copy_l2_button, clear_button]:
            component.setFont(self.ui_font)
            top_panel.add(component)
        
        column_names = [u"主域名", u"子域名", u"URL", u"一级路由", u"二级路由"]
        self.table_model = DefaultTableModel(None, column_names)
        self.log_table = JTable(self.table_model)
        
        self.log_table.setFont(self.ui_font)
        self.log_table.getTableHeader().setFont(self.ui_font)
        
        self.sorter = TableRowSorter(self.table_model)
        self.log_table.setRowSorter(self.sorter)

        scroll_pane = JScrollPane(self.log_table)
        self.main_panel.setTopComponent(top_panel)
        self.main_panel.setBottomComponent(scroll_pane)
        self.main_panel.setResizeWeight(0.1)

    def clear_results(self, event):
        self.table_model.setRowCount(0)
        print(u"All results have been cleared from the table.")

    def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo):
        if toolFlag == self._callbacks.TOOL_PROXY and messageIsRequest:
            request_info = self._helpers.analyzeRequest(messageInfo)
            url = request_info.getUrl()
            
            try:
                url_str = unicode(url.toString())
                
                parsed_url = URL(url_str)
                hostname = parsed_url.getHost()
                main_domain = get_main_domain(hostname)
                path_segments = [seg for seg in parsed_url.getPath().split('/') if seg]

                protocol = parsed_url.getProtocol()
                port = parsed_url.getPort()
                
                base_url = u"{}://{}".format(protocol, hostname)
                if port != -1:
                    base_url += u":{}".format(port)

                l1_route = u""
                l2_route = u""
                
                if len(path_segments) >= 1:
                    l1_route = base_url + u"/" + path_segments[0]
                
                if len(path_segments) >= 2:
                    l2_route = base_url + u"/" + path_segments[0] + u"/" + path_segments[1]

                row_data = [main_domain, hostname, url_str, l1_route, l2_route]
                SwingUtilities.invokeLater(lambda: self.table_model.addRow(row_data))

            except Exception as e:
                try:
                    url_for_error = url.toString()
                except:
                    url_for_error = "Unstringable URL"
                print(u"Error processing URL {}: {}".format(unicode(url_for_error), e))

    def apply_filter(self):
        from javax.swing import RowFilter
        filter_text = self.filter_field.getText()
        if not filter_text:
            self.sorter.setRowFilter(None)
        else:
            self.sorter.setRowFilter(RowFilter.regexFilter("(?i)" + filter_text, 0, 1))

    def export_csv(self, event):
        chooser = JFileChooser()
        chooser.setDialogTitle("Save as CSV")
        chooser.setFileSelectionMode(JFileChooser.FILES_ONLY)
        
        if chooser.showSaveDialog(self.main_panel) == JFileChooser.APPROVE_OPTION:
            file_path = chooser.getSelectedFile().getPath()
            if not file_path.lower().endswith(".csv"):
                file_path += ".csv"
            
            try:
                with codecs.open(file_path, 'w', 'utf-8-sig') as f:
                    header = u','.join([self.table_model.getColumnName(i) for i in range(self.table_model.getColumnCount())])
                    f.write(header + u'\n')
                    
                    for row_index in range(self.log_table.getRowCount()):
                        row_data = []
                        for col_index in range(self.log_table.getColumnCount()):
                            cell_value = unicode(self.log_table.getValueAt(row_index, col_index))
                            if u'"' in cell_value or u',' in cell_value:
                                cell_value = u'"' + cell_value.replace(u'"', u'""') + u'"'
                            row_data.append(cell_value)
                        f.write(u','.join(row_data) + u'\n')
                print(u"Data exported successfully to {}".format(file_path))
            except Exception as e:
                print(u"Error exporting to CSV: {}".format(e))

    def copy_routes(self, column_index):
        unique_routes = set()
        for row_index in range(self.log_table.getRowCount()):
            route = self.log_table.getValueAt(row_index, column_index)
            if route:
                unique_routes.add(unicode(route))
        
        clipboard_content = u"\n".join(sorted(list(unique_routes)))
        
        toolkit = Toolkit.getDefaultToolkit()
        clipboard = toolkit.getSystemClipboard()
        string_selection = StringSelection(clipboard_content)
        clipboard.setContents(string_selection, None)
        
        col_name = self.table_model.getColumnName(column_index)
        print(u"{} unique '{}' copied to clipboard.".format(len(unique_routes), col_name))

class FilterListener(DocumentListener):
    def __init__(self, extender):
        self._extender = extender
    def changedUpdate(self, e):
        self._extender.apply_filter()
    def removeUpdate(self, e):
        self._extender.apply_filter()
    def insertUpdate(self, e):
        self._extender.apply_filter()